import json
import logging
from typing import List, Dict, Any, Optional
from datetime import datetime
from app import db
from app.models.database import (
    Scan, Vulnerability, TaskStatus, TargetRange, Schedule,
    RepairSuggestion, RiskAssessment, AttackPath, ScanResult,
    vulnerability_suggestions
)
from app.services.repair_suggestion_service import repair_suggestion_service
from app.services.risk_assessment import risk_assessment_service

# 配置日志
logger = logging.getLogger(__name__)

class DatabaseService:
    """数据库服务层，提供数据持久化功能"""
    
    def __init__(self):
        self.logger = logger
    
    def create_scan_task(self, task_data: Dict[str, Any]) -> Scan:
        """创建扫描任务"""
        try:
            scan = Scan(
                target=task_data.get('target'),
                scan_depth=task_data.get('scan_depth', 2),
                scan_types=json.dumps(task_data.get('scan_types', [])),
                scheduled=task_data.get('scheduled', False),
                tags=json.dumps(task_data.get('tags', [])),
                is_public=task_data.get('is_public', False),
                user_id=task_data.get('user_id'),
                status='pending'
            )
            db.session.add(scan)
            db.session.commit()
            self.logger.info(f"Created scan task: {scan.id}")
            return scan
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to create scan task: {str(e)}")
            raise
    
    def update_scan_status(self, scan_id: str, status: str, **kwargs) -> Optional[Scan]:
        """更新扫描状态"""
        try:
            scan = Scan.query.get(scan_id)
            if not scan:
                self.logger.warning(f"Scan not found: {scan_id}")
                return None
            
            scan.status = status
            
            # 更新其他字段
            if status == 'running' and not scan.started_at:
                scan.started_at = datetime.utcnow()
            elif status == 'completed' or status == 'failed':
                scan.finished_at = datetime.utcnow()
            
            # 更新统计信息
            if 'total_hosts' in kwargs:
                scan.total_hosts = kwargs['total_hosts']
            if 'total_vulnerabilities' in kwargs:
                scan.total_vulnerabilities = kwargs['total_vulnerabilities']
            if 'high_vulnerabilities' in kwargs:
                scan.high_vulnerabilities = kwargs['high_vulnerabilities']
            if 'medium_vulnerabilities' in kwargs:
                scan.medium_vulnerabilities = kwargs['medium_vulnerabilities']
            if 'low_vulnerabilities' in kwargs:
                scan.low_vulnerabilities = kwargs['low_vulnerabilities']
            
            db.session.commit()
            self.logger.info(f"Updated scan {scan_id} status to {status}")
            return scan
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to update scan status: {str(e)}")
            raise
    
    def save_vulnerability(self, scan_id: str, vuln_data: Dict[str, Any]) -> Vulnerability:
        """保存漏洞信息"""
        try:
            vulnerability = Vulnerability(
                scan_id=scan_id,
                vuln_id=vuln_data.get('vuln_id'),
                vuln_type=vuln_data.get('vuln_type'),
                path=vuln_data.get('path'),
                payload=vuln_data.get('payload'),
                evidence=vuln_data.get('evidence'),
                raw_score=vuln_data.get('raw_score'),
                risk_score=vuln_data.get('risk_score'),
                severity=vuln_data.get('severity'),
                suggestion=vuln_data.get('suggestion')
            )
            db.session.add(vulnerability)
            db.session.flush()  # 获取ID但不提交事务
            
            # 异步生成修复建议
            self._process_vulnerability_after_save(vulnerability, vuln_data)
            
            db.session.commit()
            self.logger.info(f"Saved vulnerability: {vulnerability.id} ({vulnerability.vuln_type})")
            return vulnerability
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to save vulnerability: {str(e)}")
            raise
    
    def _process_vulnerability_after_save(self, vulnerability: Vulnerability, vuln_data: Dict[str, Any]):
        """漏洞保存后的处理：生成风险评估和修复建议"""
        try:
            # 1. 生成风险评估
            risk_data = self._generate_risk_assessment(vuln_data)
            if risk_data:
                risk_assessment = RiskAssessment(
                    scan_id=vulnerability.scan_id,
                    vulnerability_id=vulnerability.id,
                    risk_score=risk_data['risk_score'],
                    risk_level=risk_data['risk_level'],
                    exploitation_probability=risk_data.get('exploitation_probability', 0.5),
                    potential_impact=risk_data.get('potential_impact', 0.5),
                    data_sensitivity=risk_data.get('data_sensitivity', 0.5),
                    business_criticality=risk_data.get('business_criticality', 0.5),
                    feature_vector=json.dumps(risk_data.get('feature_vector', {})),
                    model_version=risk_data.get('model_version'),
                    prediction_confidence=risk_data.get('prediction_confidence', 0.9)
                )
                db.session.add(risk_assessment)
            
            # 2. 生成修复建议
            suggestion_data = self._generate_repair_suggestion(vuln_data)
            if suggestion_data:
                suggestion = RepairSuggestion(
                    title=suggestion_data['title'],
                    description=suggestion_data['description'],
                    fix_steps=json.dumps(suggestion_data.get('fix_steps', [])),
                    code_example=suggestion_data.get('code_example'),
                    best_practices=json.dumps(suggestion_data.get('best_practices', [])),
                    recommended_tools=json.dumps(suggestion_data.get('tools', [])),
                    source=suggestion_data.get('source', 'ai')
                )
                db.session.add(suggestion)
                db.session.flush()
                
                # 关联漏洞和修复建议
                vulnerability.repair_suggestions.append(suggestion)
        except Exception as e:
            self.logger.warning(f"Failed to process vulnerability after save: {str(e)}")
    
    def _generate_risk_assessment(self, vuln_data: Dict[str, Any]) -> Dict[str, Any]:
        """生成风险评估"""
        try:
            # 使用风险评估服务
            risk_result = risk_assessment_service.predict_risk(vuln_data)
            return risk_result
        except Exception as e:
            self.logger.warning(f"Failed to generate risk assessment: {str(e)}")
            # 返回默认风险评估
            return {
                'risk_score': vuln_data.get('risk_score', 0.5),
                'risk_level': vuln_data.get('severity', 'medium'),
                'source': 'fallback'
            }
    
    def _generate_repair_suggestion(self, vuln_data: Dict[str, Any]) -> Dict[str, Any]:
        """生成修复建议"""
        try:
            # 使用修复建议服务
            suggestion = repair_suggestion_service.get_repair_suggestion(vuln_data)
            return suggestion
        except Exception as e:
            self.logger.warning(f"Failed to generate repair suggestion: {str(e)}")
            return None
    
    def save_scan_result(self, scan_id: str, result_data: Dict[str, Any]) -> ScanResult:
        """保存扫描结果详情"""
        try:
            scan_result = ScanResult(
                scan_id=scan_id,
                scan_duration=result_data.get('scan_duration'),
                scan_status=result_data.get('scan_status', 'completed'),
                discovered_urls=json.dumps(result_data.get('discovered_urls', [])),
                discovered_forms=json.dumps(result_data.get('discovered_forms', [])),
                discovered_endpoints=json.dumps(result_data.get('discovered_endpoints', [])),
                security_headers=json.dumps(result_data.get('security_headers', {})),
                ssl_tls_info=json.dumps(result_data.get('ssl_tls_info', {})),
                cookie_analysis=json.dumps(result_data.get('cookie_analysis', {})),
                notes=result_data.get('notes'),
                metadata=json.dumps(result_data.get('metadata', {}))
            )
            db.session.add(scan_result)
            db.session.commit()
            self.logger.info(f"Saved scan result: {scan_result.id}")
            return scan_result
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to save scan result: {str(e)}")
            raise
    
    def save_attack_path(self, scan_id: str, path_data: Dict[str, Any]) -> AttackPath:
        """保存攻击路径"""
        try:
            attack_path = AttackPath(
                scan_id=scan_id,
                path_name=path_data.get('path_name', 'Untitled Attack Path'),
                path_score=path_data.get('path_score', 0.0),
                path_level=path_data.get('path_level', 'medium'),
                path_steps=json.dumps(path_data.get('path_steps', [])),
                path_length=len(path_data.get('path_steps', [])),
                target_host=path_data.get('target_host'),
                target_service=path_data.get('target_service'),
                target_asset=path_data.get('target_asset'),
                attack_techniques=json.dumps(path_data.get('attack_techniques', [])),
                required_privileges=path_data.get('required_privileges'),
                difficulty=path_data.get('difficulty', 'medium')
            )
            db.session.add(attack_path)
            db.session.commit()
            self.logger.info(f"Saved attack path: {attack_path.id}")
            return attack_path
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to save attack path: {str(e)}")
            raise
    
    def get_scan_by_id(self, scan_id: str) -> Optional[Scan]:
        """根据ID获取扫描任务"""
        return Scan.query.get(scan_id)
    
    def get_vulnerabilities_by_scan_id(self, scan_id: str) -> List[Vulnerability]:
        """获取扫描任务的所有漏洞"""
        return Vulnerability.query.filter_by(scan_id=scan_id).all()
    
    def get_vulnerability_count_by_severity(self, scan_id: str) -> Dict[str, int]:
        """统计扫描任务的漏洞严重程度分布"""
        vulnerabilities = self.get_vulnerabilities_by_scan_id(scan_id)
        counts = {'low': 0, 'medium': 0, 'high': 0, 'critical': 0}
        
        for vuln in vulnerabilities:
            severity = vuln.severity.lower() if vuln.severity else 'medium'
            if severity in counts:
                counts[severity] += 1
        
        return counts
    
    def get_scan_with_details(self, scan_id: str) -> Dict[str, Any]:
        """获取扫描任务的详细信息，包括漏洞、风险评估和修复建议"""
        try:
            scan = self.get_scan_by_id(scan_id)
            if not scan:
                return None
            
            # 基础信息
            scan_data = {
                'id': scan.id,
                'target': scan.target,
                'status': scan.status,
                'scan_depth': scan.scan_depth,
                'scan_types': json.loads(scan.scan_types) if scan.scan_types else [],
                'created_at': scan.created_at.isoformat(),
                'finished_at': scan.finished_at.isoformat() if scan.finished_at else None,
                'total_hosts': scan.total_hosts,
                'total_vulnerabilities': scan.total_vulnerabilities,
                'high_vulnerabilities': scan.high_vulnerabilities,
                'medium_vulnerabilities': scan.medium_vulnerabilities,
                'low_vulnerabilities': scan.low_vulnerabilities,
                'tags': json.loads(scan.tags) if scan.tags else [],
                'vulnerabilities': []
            }
            
            # 获取漏洞详情
            for vulnerability in scan.vulnerabilities:
                vuln_data = vulnerability.to_dict()
                
                # 添加风险评估信息
                risk_assessment = RiskAssessment.query.filter_by(
                    vulnerability_id=vulnerability.id
                ).first()
                if risk_assessment:
                    vuln_data['risk_assessment'] = risk_assessment.to_dict()
                
                # 添加修复建议
                suggestions = []
                for suggestion in vulnerability.repair_suggestions:
                    suggestions.append(suggestion.to_dict())
                vuln_data['repair_suggestions'] = suggestions
                
                scan_data['vulnerabilities'].append(vuln_data)
            
            # 获取攻击路径
            attack_paths = AttackPath.query.filter_by(scan_id=scan_id).all()
            scan_data['attack_paths'] = [path.to_dict() for path in attack_paths]
            
            # 获取扫描结果详情
            scan_result = ScanResult.query.filter_by(scan_id=scan_id).first()
            if scan_result:
                scan_data['scan_result'] = scan_result.to_dict()
            
            return scan_data
        except Exception as e:
            self.logger.error(f"Failed to get scan details: {str(e)}")
            return None
    
    def update_task_status(self, task_id: str, status: str, progress: int = None, 
                          result: Dict[str, Any] = None, error: str = None) -> Optional[TaskStatus]:
        """更新任务状态（用于Celery任务跟踪）"""
        try:
            task_status = TaskStatus.query.filter_by(task_id=task_id).first()
            
            if not task_status:
                # 创建新的任务状态
                task_status = TaskStatus(
                    task_id=task_id,
                    status=status,
                    progress=progress or 0
                )
                db.session.add(task_status)
            else:
                # 更新现有任务状态
                task_status.status = status
                if progress is not None:
                    task_status.progress = progress
            
            if result:
                task_status.result = json.dumps(result)
            if error:
                task_status.error = error
            
            task_status.updated_at = datetime.utcnow()
            db.session.commit()
            return task_status
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to update task status: {str(e)}")
            return None
    
    def delete_scan(self, scan_id: str) -> bool:
        """删除扫描任务及其相关数据"""
        try:
            # 开始事务
            
            # 删除关联的修复建议关联
            db.session.execute(
                vulnerability_suggestions.delete().where(
                    vulnerability_suggestions.c.vulnerability_id.in_(
                        db.session.query(Vulnerability.id).filter_by(scan_id=scan_id)
                    )
                )
            )
            
            # 删除关联的数据
            RiskAssessment.query.filter_by(scan_id=scan_id).delete()
            AttackPath.query.filter_by(scan_id=scan_id).delete()
            ScanResult.query.filter_by(scan_id=scan_id).delete()
            
            # 删除漏洞及其关联的修复建议（如果没有其他漏洞引用）
            vulnerabilities = Vulnerability.query.filter_by(scan_id=scan_id).all()
            for vuln in vulnerabilities:
                # 获取仅与此漏洞关联的修复建议
                for suggestion in vuln.repair_suggestions:
                    if suggestion.vulnerabilities.count() == 1:
                        db.session.delete(suggestion)
                db.session.delete(vuln)
            
            # 删除任务状态
            TaskStatus.query.filter_by(scan_id=scan_id).delete()
            
            # 删除扫描任务
            scan = Scan.query.get(scan_id)
            if scan:
                db.session.delete(scan)
            
            db.session.commit()
            self.logger.info(f"Deleted scan and related data: {scan_id}")
            return True
        except Exception as e:
            db.session.rollback()
            self.logger.error(f"Failed to delete scan: {str(e)}")
            return False
    
    def get_recent_scans(self, limit: int = 10, user_id: str = None) -> List[Scan]:
        """获取最近的扫描任务"""
        query = Scan.query.order_by(Scan.created_at.desc())
        
        if user_id:
            query = query.filter_by(user_id=user_id)
        
        return query.limit(limit).all()
    
    def get_vulnerability_trends(self, days: int = 30) -> List[Dict[str, Any]]:
        """获取漏洞趋势统计"""
        from datetime import timedelta
        
        end_date = datetime.utcnow()
        start_date = end_date - timedelta(days=days)
        
        # 获取指定时间范围内的扫描
        scans = Scan.query.filter(
            Scan.created_at >= start_date,
            Scan.created_at <= end_date
        ).all()
        
        # 按日期分组统计
        trends = []
        for scan in scans:
            vuln_count = len(scan.vulnerabilities)
            trends.append({
                'date': scan.created_at.date().isoformat(),
                'total_vulnerabilities': vuln_count,
                'high': sum(1 for v in scan.vulnerabilities if v.severity == 'high'),
                'medium': sum(1 for v in scan.vulnerabilities if v.severity == 'medium'),
                'low': sum(1 for v in scan.vulnerabilities if v.severity == 'low')
            })
        
        return trends

# 全局数据库服务实例
database_service = DatabaseService()

# 便捷函数
def create_scan_task(task_data: Dict[str, Any]) -> Scan:
    """创建扫描任务（便捷函数）"""
    return database_service.create_scan_task(task_data)

def update_scan_status(scan_id: str, status: str, **kwargs) -> Optional[Scan]:
    """更新扫描状态（便捷函数）"""
    return database_service.update_scan_status(scan_id, status, **kwargs)

def save_vulnerability(scan_id: str, vuln_data: Dict[str, Any]) -> Vulnerability:
    """保存漏洞信息（便捷函数）"""
    return database_service.save_vulnerability(scan_id, vuln_data)

def save_scan_result(scan_id: str, result_data: Dict[str, Any]) -> ScanResult:
    """保存扫描结果（便捷函数）"""
    return database_service.save_scan_result(scan_id, result_data)

def get_scan_with_details(scan_id: str) -> Dict[str, Any]:
    """获取扫描详情（便捷函数）"""
    return database_service.get_scan_with_details(scan_id)